"""
12_community_detection.py
=========================
Community Detection Algorithms with GDS

This module covers:
1. Louvain community detection
2. Label Propagation
3. Weakly Connected Components
4. Strongly Connected Components
5. K-1 Coloring
6. Triangle Count

Author: Neo4j Tutorial
Date: January 2025
"""

from typing import Dict, List, Any, Optional
from neo4j import Driver
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.utils.connection import get_driver


class CommunityDetection:
    """
    Implements community detection algorithms using Neo4j GDS.
    
    Community detection finds groups of densely connected nodes.
    """
    
    def __init__(self, driver: Driver):
        self.driver = driver
        self.graph_name = "community_graph"
    
    def create_graph_projection(self) -> Dict[str, Any]:
        """Create graph projection for community detection."""
        drop_query = f"CALL gds.graph.drop('{self.graph_name}', false)"
        
        create_query = """
        CALL gds.graph.project(
            $graph_name,
            ['Patient', 'Disease', 'Medication'],
            {
                DIAGNOSED_WITH: {orientation: 'UNDIRECTED'},
                TAKES: {orientation: 'UNDIRECTED'},
                TREATS: {orientation: 'UNDIRECTED'}
            }
        )
        YIELD graphName, nodeCount, relationshipCount
        RETURN graphName, nodeCount, relationshipCount
        """
        
        with self.driver.session() as session:
            try:
                session.run(drop_query)
            except:
                pass
            
            result = session.run(create_query, graph_name=self.graph_name)
            record = result.single()
            return record.data() if record else None
    
    # =========================================================
    # LOUVAIN COMMUNITY DETECTION
    # =========================================================
    
    def louvain(
        self,
        max_levels: int = 10,
        max_iterations: int = 10,
        tolerance: float = 0.0001,
        include_intermediate: bool = False
    ) -> Dict[str, Any]:
        """
        Run Louvain community detection algorithm.
        
        Louvain is a hierarchical clustering algorithm that optimizes
        modularity. It's fast and works well on large graphs.
        
        Args:
            max_levels: Maximum hierarchy levels
            max_iterations: Iterations per level
            tolerance: Modularity gain threshold
            include_intermediate: Include intermediate communities
            
        Returns:
            Community statistics and assignments
            
        Use Cases:
            - Group patients with similar conditions
            - Find medication clusters
            - Identify disease families
        """
        query = """
        CALL gds.louvain.stream($graph_name, {
            maxLevels: $max_levels,
            maxIterations: $max_iter,
            tolerance: $tol,
            includeIntermediateCommunities: $intermediate
        })
        YIELD nodeId, communityId, intermediateCommunityIds
        WITH gds.util.asNode(nodeId) AS node, communityId, intermediateCommunityIds
        RETURN 
            node.id AS id,
            node.name AS name,
            labels(node)[0] AS label,
            communityId,
            intermediateCommunityIds
        ORDER BY communityId, name
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                graph_name=self.graph_name,
                max_levels=max_levels,
                max_iter=max_iterations,
                tol=tolerance,
                intermediate=include_intermediate
            )
            records = [record.data() for record in result]
            
            # Calculate statistics
            communities = {}
            for r in records:
                cid = r['communityId']
                if cid not in communities:
                    communities[cid] = []
                communities[cid].append(r)
            
            return {
                "total_communities": len(communities),
                "total_nodes": len(records),
                "community_sizes": {k: len(v) for k, v in communities.items()},
                "nodes": records[:50]  # Sample
            }
    
    def louvain_write(self, property_name: str = "community") -> Dict[str, Any]:
        """
        Run Louvain and write community IDs to nodes.
        
        Args:
            property_name: Property name to write
            
        Returns:
            Write statistics
        """
        query = """
        CALL gds.louvain.write($graph_name, {
            writeProperty: $prop
        })
        YIELD communityCount, modularity, nodePropertiesWritten
        RETURN communityCount, modularity, nodePropertiesWritten
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                graph_name=self.graph_name,
                prop=property_name
            )
            return result.single().data()
    
    # =========================================================
    # LABEL PROPAGATION
    # =========================================================
    
    def label_propagation(
        self,
        max_iterations: int = 10
    ) -> Dict[str, Any]:
        """
        Run Label Propagation algorithm.
        
        Fast, semi-supervised community detection where nodes
        adopt the most common label among neighbors.
        
        Args:
            max_iterations: Maximum iterations
            
        Returns:
            Community assignments
            
        Use Cases:
            - Quick community overview
            - Semi-supervised clustering with seed labels
        """
        query = """
        CALL gds.labelPropagation.stream($graph_name, {
            maxIterations: $max_iter
        })
        YIELD nodeId, communityId
        WITH gds.util.asNode(nodeId) AS node, communityId
        RETURN 
            node.id AS id,
            node.name AS name,
            labels(node)[0] AS label,
            communityId
        ORDER BY communityId, name
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                graph_name=self.graph_name,
                max_iter=max_iterations
            )
            records = [record.data() for record in result]
            
            # Statistics
            communities = {}
            for r in records:
                cid = r['communityId']
                if cid not in communities:
                    communities[cid] = []
                communities[cid].append(r)
            
            return {
                "total_communities": len(communities),
                "total_nodes": len(records),
                "community_sizes": {k: len(v) for k, v in communities.items()},
                "nodes": records[:50]
            }
    
    # =========================================================
    # WEAKLY CONNECTED COMPONENTS
    # =========================================================
    
    def weakly_connected_components(self) -> Dict[str, Any]:
        """
        Find Weakly Connected Components.
        
        Groups nodes that can reach each other ignoring edge direction.
        
        Returns:
            Component assignments
            
        Use Cases:
            - Find disconnected subgraphs
            - Identify isolated node clusters
        """
        query = """
        CALL gds.wcc.stream($graph_name)
        YIELD nodeId, componentId
        WITH gds.util.asNode(nodeId) AS node, componentId
        RETURN 
            node.id AS id,
            node.name AS name,
            labels(node)[0] AS label,
            componentId
        ORDER BY componentId, name
        """
        
        with self.driver.session() as session:
            result = session.run(query, graph_name=self.graph_name)
            records = [record.data() for record in result]
            
            components = {}
            for r in records:
                cid = r['componentId']
                if cid not in components:
                    components[cid] = []
                components[cid].append(r)
            
            return {
                "total_components": len(components),
                "total_nodes": len(records),
                "component_sizes": {k: len(v) for k, v in components.items()},
                "largest_component": max(len(v) for v in components.values()) if components else 0,
                "nodes": records[:50]
            }
    
    # =========================================================
    # TRIANGLE COUNT
    # =========================================================
    
    def triangle_count(self) -> List[Dict[str, Any]]:
        """
        Count triangles for each node.
        
        Triangles indicate clustering/cohesion. High triangle count
        means the node's neighbors are also connected to each other.
        
        Returns:
            Nodes with triangle counts
            
        Use Cases:
            - Measure local clustering
            - Find tightly-knit groups
        """
        query = """
        CALL gds.triangleCount.stream($graph_name)
        YIELD nodeId, triangleCount
        WITH gds.util.asNode(nodeId) AS node, triangleCount
        WHERE triangleCount > 0
        RETURN 
            node.id AS id,
            node.name AS name,
            labels(node)[0] AS label,
            triangleCount
        ORDER BY triangleCount DESC
        LIMIT 20
        """
        
        with self.driver.session() as session:
            result = session.run(query, graph_name=self.graph_name)
            return [record.data() for record in result]
    
    def local_clustering_coefficient(self) -> List[Dict[str, Any]]:
        """
        Calculate local clustering coefficient.
        
        Measures how close a node's neighbors are to being a clique.
        
        Returns:
            Nodes with clustering coefficients
        """
        query = """
        CALL gds.localClusteringCoefficient.stream($graph_name)
        YIELD nodeId, localClusteringCoefficient
        WITH gds.util.asNode(nodeId) AS node, localClusteringCoefficient
        WHERE localClusteringCoefficient > 0
        RETURN 
            node.id AS id,
            node.name AS name,
            labels(node)[0] AS label,
            localClusteringCoefficient AS coefficient
        ORDER BY localClusteringCoefficient DESC
        LIMIT 20
        """
        
        with self.driver.session() as session:
            result = session.run(query, graph_name=self.graph_name)
            return [record.data() for record in result]
    
    # =========================================================
    # COMMUNITY ANALYSIS
    # =========================================================
    
    def analyze_community(self, community_id: int) -> Dict[str, Any]:
        """
        Analyze a specific community.
        
        Args:
            community_id: Community ID to analyze
            
        Returns:
            Community details
        """
        query = """
        MATCH (n)
        WHERE n.community = $community_id
        WITH collect(n) AS members
        UNWIND members AS m
        OPTIONAL MATCH (m)-[r]-(other)
        WHERE other IN members
        WITH members, count(DISTINCT r) AS internal_edges
        UNWIND members AS m
        OPTIONAL MATCH (m)-[r]-(external)
        WHERE NOT external IN members
        WITH members, internal_edges, count(DISTINCT r) AS external_edges
        RETURN 
            size(members) AS size,
            [m IN members | {name: m.name, label: labels(m)[0]}][..10] AS sample_members,
            internal_edges,
            external_edges,
            CASE WHEN internal_edges + external_edges > 0 
                THEN toFloat(internal_edges) / (internal_edges + external_edges) 
                ELSE 0 
            END AS density
        """
        
        with self.driver.session() as session:
            result = session.run(query, community_id=community_id)
            record = result.single()
            return record.data() if record else None
    
    def cleanup(self):
        """Drop the graph projection."""
        query = f"CALL gds.graph.drop('{self.graph_name}', false)"
        with self.driver.session() as session:
            try:
                session.run(query)
            except:
                pass


# ============================================================
# DEMONSTRATION
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("Community Detection - Demonstration")
    print("=" * 60)
    
    driver = get_driver()
    community = CommunityDetection(driver)
    
    # Create projection
    print("\n📌 Creating Graph Projection")
    print("-" * 40)
    proj = community.create_graph_projection()
    if proj:
        print(f"   Nodes: {proj['nodeCount']}")
        print(f"   Relationships: {proj['relationshipCount']}")
    
    # Louvain
    print("\n📌 Louvain Community Detection")
    print("-" * 40)
    try:
        louvain_result = community.louvain()
        print(f"   Found {louvain_result['total_communities']} communities")
        print(f"   Total nodes: {louvain_result['total_nodes']}")
        print("   Community sizes:")
        for cid, size in sorted(louvain_result['community_sizes'].items())[:5]:
            print(f"      Community {cid}: {size} nodes")
    except Exception as e:
        print(f"   Error: {e}")
    
    # Label Propagation
    print("\n📌 Label Propagation")
    print("-" * 40)
    try:
        lp_result = community.label_propagation()
        print(f"   Found {lp_result['total_communities']} communities")
    except Exception as e:
        print(f"   Error: {e}")
    
    # WCC
    print("\n📌 Weakly Connected Components")
    print("-" * 40)
    try:
        wcc_result = community.weakly_connected_components()
        print(f"   Found {wcc_result['total_components']} components")
        print(f"   Largest component: {wcc_result['largest_component']} nodes")
    except Exception as e:
        print(f"   Error: {e}")
    
    # Triangle Count
    print("\n📌 Triangle Count")
    print("-" * 40)
    try:
        triangles = community.triangle_count()
        for t in triangles[:5]:
            print(f"   {t['name']}: {t['triangleCount']} triangles")
    except Exception as e:
        print(f"   Error: {e}")
    
    # Cleanup
    community.cleanup()
    driver.close()
    
    print("\n" + "=" * 60)
    print("✅ Community detection demonstration complete!")
    print("=" * 60)
